import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler, OneHotEncoder
from StreamLearn.Dataset.StreamDataset import StreamDataset


class AdultDataset(StreamDataset):
    def __init__(self, args, file=None):
        self.name = 'adult'
        self.args = args
        self.val_size = args.val_size
        if args.dataset_mode not in ['balance', 'imbalance']:
            raise ValueError("Unknown mode {}.".format(args.dataset_mode))
        self.dataset_mode = args.dataset_mode        # ['balance', 'imbalance']
        if file is None:
            self.train_X, self.train_Y, self.test_X, self.test_Y = get_data(self.dataset_mode, self.val_size)
        else:
            self.train_X, self.train_Y, self.test_X, self.test_Y = get_data(self.dataset_mode, self.val_size, file)
        self.m = len(self.train_X)
        self.size = self.train_X[0].shape[1]
        if self.dataset_mode == 'balance':
            self.n = args.T
        elif self.dataset_mode == 'imbalance':
            self.train_X_backup = self.train_X.copy()
            self.train_Y_backup = self.train_Y.copy()
            self.n = np.array([len(i) for i in self.train_X])

    def get_m(self):
        return self.m

    def get_w_size(self):
        return self.size

    def reset(self):
        self.train_X, self.train_Y, self.test_X, self.test_Y = get_data(self.dataset_mode, self.val_size)
        if self.dataset_mode == 'imbalance':
            self.train_X_backup = self.train_X.copy()
            self.train_Y_backup = self.train_Y.copy()
            self.n = np.array([len(i) for i in self.train_X])

    def data_reset(self):
        if self.dataset_mode == 'imbalance':
            self.train_X = self.train_X_backup.copy()
            self.train_Y = self.train_Y_backup.copy()
        elif self.dataset_mode == 'balance':
            pass

    def get_n(self):
        return self.n

    def get_testdata(self, k):
        """
        for validation
        :param k: kth distribution
        :return: num * size+1 (x+y)
        """
        res = np.hstack((self.test_X[k], np.reshape(self.test_Y[k], (-1, 1))))
        return res

    def sample_m(self):
        """
        Extract 1 sample for each distribution, and return an array with size m * (size+1)
        """
        res = np.zeros((self.m, self.size + 1))
        if self.dataset_mode == 'balance':
            for i in range(self.m):
                random_number = np.random.choice(self.train_X[i].shape[0], 1)
                res[i, :-1] = self.train_X[i][random_number, :]
                res[i, -1] = self.train_Y[i][random_number]
        elif self.dataset_mode == 'imbalance':
            for i in range(self.m):
                random_rows = np.random.choice(self.train_X[i].shape[0], 1, replace=False)
                res[i, :-1] = self.train_X[i][random_rows, :]
                res[i, -1] = self.train_Y[i][random_rows]
                self.train_X[i] = np.delete(self.train_X[i], random_rows, axis=0)
                self.train_Y[i] = np.delete(self.train_Y[i], random_rows, axis=0)
        return res

    def sample_single(self, k, num):
        """
        :param k: kth distribution
        :param num: number of samples  (for adult is None
        :return: num * size+1 (x+y)
        """
        res = np.zeros((num, self.size + 1))
        if self.dataset_mode == 'balance':
            random_number = np.random.choice(self.train_X[k].shape[0], num, replace=True)
            res[:, :-1] = self.train_X[k][random_number, :]
            res[:, -1] = self.train_Y[k][random_number]
        elif self.dataset_mode == 'imbalance':
            random_rows = np.random.choice(self.train_X[k].shape[0], num, replace=False)
            res[:, :-1] = self.train_X[k][random_rows, :]
            res[:, -1] = self.train_Y[k][random_rows]
            self.train_X[k] = np.delete(self.train_X[k], random_rows, axis=0)
            self.train_Y[k] = np.delete(self.train_Y[k], random_rows, axis=0)
        return res

    def sample_pr_avail(self):
        """
        :return: some samples of size+1
        """
        counts = np.floor(self.n / np.min(self.n)).astype(int)
        res = []
        for i in range(self.m):
            random_rows = np.random.choice(self.train_X[i].shape[0], counts[i], replace=False)
            x = self.train_X[i][random_rows, :]
            out = self.train_Y[i][random_rows]
            self.train_X[i] = np.delete(self.train_X[i], random_rows, axis=0)
            self.train_Y[i] = np.delete(self.train_Y[i], random_rows, axis=0)
            res.append(np.hstack((i * np.ones((counts[i], 1)), x, out.reshape(len(out), -1))))
        return np.vstack([i for i in res])


def get_data(mode, test_num, file=None):
    train_data = pd.read_csv('Dataset/adult_train.csv')
    test_data = pd.read_csv('Dataset/adult_test.csv')

    data = pd.concat([train_data, test_data], ignore_index=True)

    data.loc[(data['race'] != 'White') & (data['race'] != 'Black'), 'race'] = 'Other'

    for col in data.columns:
        data = data[~data[col].isin(['?'])]
    data = data.reset_index(drop=True)

    group_nums = []
    for race in ['White', 'Black', 'Other']:
        for sex in ['Male', 'Female']:
            group_nums.append(len(data[(data['race'] == race) & (data['sex'] == sex)]))

    # one-hot
    object_columns = data.dtypes == "object"
    object_columns = list(object_columns[object_columns].index)
    int_columns = data.dtypes == "int64"
    int_columns = list(int_columns[int_columns].index)
    for col_name in object_columns:
        values = np.array(data[col_name])
        onehot_encoder = OneHotEncoder(sparse_output=False)
        values = values.reshape(len(values), 1)
        onehot_matrix = onehot_encoder.fit_transform(values)
        data.drop([col_name], axis=1, inplace=True)
        for i in range(onehot_matrix.shape[1]):
            data.insert(0, col_name + "_" + str(i), value=onehot_matrix[:, i])

    for col_name in int_columns:
        Scaler = MinMaxScaler(feature_range=(0, 1))
        col_value = np.array(data[col_name]).reshape(-1, 1)
        new_col = Scaler.fit_transform(col_value)
        data[col_name] = new_col

    bias_column = pd.DataFrame({'bias': 1}, index=data.index)
    data = pd.concat([data, bias_column], axis=1)

    adult_g1 = data[(data['race_0'] == 1) & (data['sex_0'] == 1)]
    adult_g2 = data[(data['race_0'] == 1) & (data['sex_1'] == 1)]
    adult_g3 = data[(data['race_1'] == 1) & (data['sex_0'] == 1)]
    adult_g4 = data[(data['race_1'] == 1) & (data['sex_1'] == 1)]
    adult_g5 = data[(data['race_2'] == 1) & (data['sex_0'] == 1)]
    adult_g6 = data[(data['race_2'] == 1) & (data['sex_1'] == 1)]
    adult_g = [adult_g1, adult_g2, adult_g3, adult_g4, adult_g5, adult_g6]
    key = lambda x: len(x)
    adult_g = sorted(adult_g, key=key, reverse=True)

    if file is not None:
        dataset = np.load(file)
        indexes = dataset['index']
        train_adult_g, test_adult_g = [], []
        for i, _ in enumerate(adult_g):
            index = indexes[i]
            test_adult_g.append(adult_g[i].iloc[index])
            train_adult_g.append(adult_g[i].iloc[np.array(list(set(range(adult_g[i].shape[0])) - set(index)))])

        train_X, train_Y, test_X, test_Y = [], [], [], []
        for i in range(6):
            train_Y1 = (train_adult_g[i].iloc[:, 0] == 1)
            train_Y1 = np.array(train_Y1, dtype=int)
            train_adult_g[i] = train_adult_g[i].loc[:, train_adult_g[i].columns.difference(["output_1", "output_0"])]
            train_X1 = train_adult_g[i].values
            train_X.append(train_X1)
            train_Y.append(train_Y1)

            test_Y1 = (test_adult_g[i].iloc[:, 0] == 1)
            test_Y1 = np.array(test_Y1, dtype=int)
            test_adult_g[i] = test_adult_g[i].loc[:, test_adult_g[i].columns.difference(["output_1", "output_0"])]
            test_X1 = test_adult_g[i].values
            test_X.append(test_X1)
            test_Y.append(test_Y1)

        return train_X, train_Y, test_X, test_Y

    if mode == 'balance':
        train_adult_g, test_adult_g = [], []
        for i, data in enumerate(adult_g):
            train_adult_g.append(adult_g[i])
            sample_index = np.random.choice(len(data), test_num, replace=True)
            test_adult_g.append(adult_g[i].iloc[sample_index])

        train_X, train_Y, test_X, test_Y = [], [], [], []

        for i in range(6):
            train_Y1 = (train_adult_g[i].iloc[:, 0] == 1)
            train_Y1 = np.array(train_Y1, dtype=int)
            train_adult_g[i] = train_adult_g[i].loc[:, train_adult_g[i].columns.difference(["output_1", "output_0"])]
            train_X1 = train_adult_g[i].values
            train_X.append(train_X1)
            train_Y.append(train_Y1)

            test_Y1 = (test_adult_g[i].iloc[:, 0] == 1)
            test_Y1 = np.array(test_Y1, dtype=int)
            test_adult_g[i] = test_adult_g[i].loc[:, test_adult_g[i].columns.difference(["output_1", "output_0"])]
            test_X1 = test_adult_g[i].values
            test_X.append(test_X1)
            test_Y.append(test_Y1)
        return train_X, train_Y, test_X, test_Y

    elif mode == 'imbalance':
        num = 364
        train_adult_g, test_adult_g = [], []
        for i, _ in enumerate(adult_g):
            index = np.random.choice(adult_g[i].shape[0], size=num, replace=False)
            test_adult_g.append(adult_g[i].iloc[index])
            train_adult_g.append(adult_g[i].iloc[np.array(list(set(range(adult_g[i].shape[0])) - set(index)))])

        train_X, train_Y, test_X, test_Y = [], [], [], []
        for i in range(6):
            train_Y1 = (train_adult_g[i].iloc[:, 0] == 1)
            train_Y1 = np.array(train_Y1, dtype=int)
            train_adult_g[i] = train_adult_g[i].loc[:, train_adult_g[i].columns.difference(["output_1", "output_0"])]
            train_X1 = train_adult_g[i].values
            train_X.append(train_X1)
            train_Y.append(train_Y1)

            test_Y1 = (test_adult_g[i].iloc[:, 0] == 1)
            test_Y1 = np.array(test_Y1, dtype=int)
            test_adult_g[i] = test_adult_g[i].loc[:, test_adult_g[i].columns.difference(["output_1", "output_0"])]
            test_X1 = test_adult_g[i].values
            test_X.append(test_X1)
            test_Y.append(test_Y1)

        return train_X, train_Y, test_X, test_Y
